/*
 * MAtoms.cpp
 *
 *  Created on: Aug 7, 2013
 *      Author: marchi
 */

#include "MAtoms.h"


template <typename T, typename U>
class create_map{
private:
    map<T, U> m_map;
public:
    create_map(const T& key, const U& val){m_map[key] = val;}

    create_map<T, U>& operator()(const T& key, const U& val) {
    	m_map[key] = val;
    	return *this;
    }

    operator map<T, U>(){return m_map;}
};

template <typename T>
class create_vector{
private:
    vector<T> m_vec;
public:
    create_vector(const T& data){m_vec.push_back(data);}

    create_vector<T>& operator()(const T& data) {
    	m_vec.push_back(data);
    	return *this;
    }

    operator vector<T>(){return m_vec;}
};

int MAtoms::calls=0;
vector<string> * MAtoms::ResList0=NULL;
vector<int> * MAtoms::ResIndx0=NULL;

map<string,double> MAtoms::MMass=create_map<string,double>
("H",1.008)("C",12.011)("N",14.0067)("O",15.9994)("S",32.06)("P",30.974)
("FE",55.847)("CA",40.080)("ZN",65.37)("NA",22.98977)("MG",24.305);
string MAtoms::MIons="CA NA SOD CAL ZN";

class PBCvect{
	static vector<Dvect> * nxyz;
	const int M;
public:
	PBCvect(): M(3){
		nxyz=new vector<Dvect>;
		for(int o=0;o<M;o++){
			double nx=static_cast<double>(o-1);
			for(int p=0;p<M;p++){
				double ny=static_cast<double>(p-1);
				for(int q=0;q<M;q++){
					double nz=static_cast<double>(q-1);
					nxyz->push_back(Dvect(nx,ny,nz));
				}
			}
		}
	}
	static vector<Dvect> & getVec(){return *nxyz;}
};
struct op_iComp {
  bool operator() (double i,double j) { return (i>j);}
};

class Jacob{
	double ** iner;
	double ** ei;
	double * Im;
	double sigma;
	static Dvect Im_s;
	double aug(double p){return p*(1+sigma);}
	void  eigsort(Dvect  & d, Matrix  * v=NULL){
		int k;
		int n=DIM;
		for (int i=0;i<n-1;i++) {
			double p=d[k=i];
			for (int j=i;j<n;j++)
				if (d[j] >= aug(p)) p=d[k=j];
			if (k != i) {
				d[k]=d[i];
				d[i]=p;
				if (v != NULL){
					for (int j=0;j<n;j++) {
						p=(*v)[j][i];
						(*v)[j][i]=(*v)[j][k];
						(*v)[j][k]=p;
					}
				}
			}
		}
	}
public:
	Jacob():  sigma(0.1), iner(NULL), ei(NULL),Im(NULL){};
	~Jacob(){
		for(int o=0;o<DIM;o++) delete [] iner[o];
		delete [] iner;
		for(int o=0;o<DIM;o++) delete [] ei[o];
		delete [] ei;
		delete [] Im;
	}
	void operator()(Matrix & iner0, Dvect & Im0, Matrix & ei0, int & nrot){
		if(!iner) {
			iner=new double * [DIM];
			for(int o=0;o<DIM;o++) iner[o]=new double [DIM];
		}
		if(!ei) {
			ei=new double * [DIM];
			for(int o=0;o<DIM;o++) ei[o]=new double [DIM];
		}
		if(!Im) Im=new double [DIM];
		for(int o=0;o<DIM;o++){
			for(int p=0;p<DIM;p++){
				iner[o][p]=iner0[o][p];
			}
		}
		jacobi(iner,DIM,Im,ei,&nrot);

		for(int o=0;o<DIM;o++){
			Im0[o]=Im[o];
			for(int p=0;p<DIM;p++){
				ei0[o][p]=ei[o][p];
			}
		}

		eigsort(Im0,&ei0);

	}

};
Dvect Jacob::Im_s=0.0;

vector<Dvect> * PBCvect::nxyz=NULL;

class CellComp{
	Dvect Xd;
	Matrix co;
public:
	CellComp(Dvect& x, Dvect & y, Matrix & c): Xd(x-y), co(c){};
	CellComp(Dvect& x, Matrix & c): Xd(x), co(c){};
	bool operator()(const Dvect & x, const Dvect & y){
		Dvect x1=Xd+x;
		Dvect x2=Xd+y;
		Dvect xc1=co*x1;
		Dvect xc2=co*x2;
		return xc1.Norm() < xc2.Norm();
	}
};

MAtoms::MAtoms(): Rg_count(0){
	// TODO Auto-generated constructor stub
}
void MAtoms::setList(vector<string> & L){
	if(L.empty()) return;
	ResList0=new vector<string>(L);
	ResIndx0=new vector<int>;
	int ia=0,ib=0;
	for(size_t i=0;i<CurrentPDB.size();i++){
		string Curr=CurrentPDB[i];
		if(Curr.find("ATOM") == 0 || Curr.find("HETATM") == 0 ) {
			string sub1=CurrentPDB[i].substr(11,5);
			string sub2=CurrentPDB[i].substr(16,4);
			sub1.erase(remove_if(sub1.begin(),sub1.end(),::isspace),sub1.end());
			sub2.erase(remove_if(sub2.begin(),sub2.end(),::isspace),sub2.end());
			if(MIons.find(sub2) == string::npos) mass.push_back(MMass[sub1.substr(0,1)]);
			else {
				mass.push_back(MMass[sub1.substr(0,2)]);
			}
			if(!ResList0 || find(ResList0->begin(),ResList0->end(),sub2) != ResList0->end()) {
				ResIndx0->push_back(ia);
			}
			ia++;
		}
	}
}
void MAtoms::Gyro(){
	Rg_i=0.0;
	Gyration::setTime(time_c);
	Dvect cm=0.0,cmG=0.0;
	double unit_nmm2=1.0/(unit_nm*unit_nm);
	int ntot=static_cast<int>(ResIndx0->size());

	double tmass=0.0;
	for(unsigned int n=0;n<ResIndx0->size();n++){
		int nn=(*ResIndx0)[n];
		for(int o=0;o<DIM;o++) cm[o]+=x[nn][o]*mass[nn];
		for(int o=0;o<DIM;o++) cmG[o]+=x[nn][o];
		tmass+=mass[nn];
	}
	cm/=tmass;
	cmG/=static_cast<int> (ntot);
	Dvect x0;
	Dvect Gm,Im,axis;
	Matrix Giner=0.0;
	Matrix ei,ei2=0.0;
	double MyRg=0;
	int oa=0;
	for(unsigned int n=0;n<ResIndx0->size();n++){
		int nn=(*ResIndx0)[n];
		for(int o=0;o<DIM;o++) x0[o]=cmG[o]-x[ResIndx0->at(n)][o];
		Giner+=x0%x0*unit_nmm2;
	}
	Giner/=static_cast<double> (ntot);
	int nrot;
	Jacob()(Giner,Gm,ei,nrot);
	for(int o=0;o<DIM;o++) MyRg+=Gm[o];
	Matrix Inertia=0.0;
	for(unsigned int n=0;n<ResIndx0->size();n++){
		int nn=(*ResIndx0)[n];
		for(int o=0;o<DIM;o++) x0[o]=cm[o]-x[nn][o];

		Inertia[XX][XX]+=(sqr(x0[YY])+sqr(x0[ZZ]))*mass[nn];
		Inertia[YY][YY]+=(sqr(x0[ZZ])+sqr(x0[XX]))*mass[nn];
		Inertia[ZZ][ZZ]+=(sqr(x0[XX])+sqr(x0[YY]))*mass[nn];
		Inertia[XX][YY]-=(x0[XX]*x0[YY])*mass[nn];
		Inertia[XX][ZZ]-=(x0[XX]*x0[ZZ])*mass[nn];
		Inertia[YY][ZZ]-=(x0[YY]*x0[ZZ])*mass[nn];
	}
	Inertia[YY][XX]=Inertia[XX][YY];
	Inertia[ZZ][XX]=Inertia[XX][ZZ];
	Inertia[ZZ][YY]=Inertia[YY][ZZ];
	Jacob()(Inertia,Im,ei2,nrot);
	Im*=unit_nmm2;
	double fact=5.0/tmass/2.0;
	axis[XX]=fact*(Im[YY]+Im[ZZ]-Im[XX]);
	axis[YY]=fact*(Im[XX]+Im[ZZ]-Im[YY]);
	axis[ZZ]=fact*(Im[YY]+Im[XX]-Im[ZZ]);
	Rg_i(MyRg,Im,Gm,axis);
	Rg_stat+=Rg_i;
	Rg_count++;

}
Dvect MAtoms::CM(){
	Dvect cm=0.0;
	double tmass=0.0;
	for(unsigned int n=0;n<ResIndx0->size();n++){
		int nn=(*ResIndx0)[n];
		for(int o=0;o<DIM;o++) cm[o]+=x[nn][o]*mass[nn];
		tmass+=mass[nn];
	}
	cm/=tmass;
	return cm;
}
Dvect MAtoms::CMCenter(){
	Dvect cm=0.0;
	double tmass=0.0;
	for(unsigned int n=0;n < MyRes.size();n++){
		for(unsigned int m=0;m < MyRes[n].size();m++){
			int nn=MyRes[n][m];
			for(int o=0;o<DIM;o++) cm[o]+=x[nn][o]*mass[nn];
			tmass+=mass[nn];
		}
	}
	cm/=tmass;
	return cm;
}
Matrix MAtoms::InertiaRotCenter(){
	Dvect x0,cm;
	cm=CMCenter();

	double unit_nmm2=1.0/(unit_nm*unit_nm);
	Matrix Inertia=0.0,ei,eiT;
	Dvect Im;
	int nrot;

	for(unsigned int n=0;n < MyRes.size();n++){
		for(unsigned int m=0;m < MyRes[n].size();m++){
			int nn=MyRes[n][m];
			for(int o=0;o<DIM;o++) x0[o]=x[nn][o]-cm[o];
			Inertia[XX][XX]+=(sqr(x0[YY])+sqr(x0[ZZ]))*mass[nn];
			Inertia[YY][YY]+=(sqr(x0[ZZ])+sqr(x0[XX]))*mass[nn];
			Inertia[ZZ][ZZ]+=(sqr(x0[XX])+sqr(x0[YY]))*mass[nn];
			Inertia[XX][YY]-=(x0[XX]*x0[YY])*mass[nn];
			Inertia[XX][ZZ]-=(x0[XX]*x0[ZZ])*mass[nn];
			Inertia[YY][ZZ]-=(x0[YY]*x0[ZZ])*mass[nn];
		}
	}
	Inertia[YY][XX]=Inertia[XX][YY];
	Inertia[ZZ][XX]=Inertia[XX][ZZ];
	Inertia[ZZ][YY]=Inertia[YY][ZZ];
	Jacob()(Inertia,Im,ei,nrot);
	eiT=ei.Transpose();
	return eiT;
}

Matrix MAtoms::InertiaRot(){
	Dvect cm=0.0,x0;
	double tmass=0.0;
	for(unsigned int n=0;n<ResIndx0->size();n++){
		int nn=(*ResIndx0)[n];
		for(int o=0;o<DIM;o++) cm[o]+=x[nn][o]*mass[nn];
		tmass+=mass[nn];
	}
	cm/=tmass;

	double unit_nmm2=1.0/(unit_nm*unit_nm);
	Matrix Inertia=0.0,ei,eiT;
	Dvect Im;
	int nrot;
	for(unsigned int n=0;n<ResIndx0->size();n++){
		int nn=(*ResIndx0)[n];
		for(int o=0;o<DIM;o++) x0[o]=x[nn][o]-cm[o];

		Inertia[XX][XX]+=(sqr(x0[YY])+sqr(x0[ZZ]))*mass[nn];
		Inertia[YY][YY]+=(sqr(x0[ZZ])+sqr(x0[XX]))*mass[nn];
		Inertia[ZZ][ZZ]+=(sqr(x0[XX])+sqr(x0[YY]))*mass[nn];
		Inertia[XX][YY]-=(x0[XX]*x0[YY])*mass[nn];
		Inertia[XX][ZZ]-=(x0[XX]*x0[ZZ])*mass[nn];
		Inertia[YY][ZZ]-=(x0[YY]*x0[ZZ])*mass[nn];
	}
	Inertia[YY][XX]=Inertia[XX][YY];
	Inertia[ZZ][XX]=Inertia[XX][ZZ];
	Inertia[ZZ][YY]=Inertia[YY][ZZ];
	Jacob()(Inertia,Im,ei,nrot);
	eiT=ei.Transpose();
	return eiT;
}
string MAtoms::gTotResList(){
	stringstream ss;
	for(vector<string>::iterator it=ResList0->begin();it != ResList0->end();it++)
		ss<< *it + " " ;
	return ss.str();
}
void MAtoms::Reconstruct(const string & y, Topol & MyTop){
	try{
		if(!MyTop.CheckResidue(y)) throw string(" Reference residue " + y + " does not exist in the system. Abort");
	} catch(const string & s){
		cout << s << endl;
		exit(1);
	}
	if(!calls++) {PBCvect dummy;}
	vector<Dvect> nxyz=PBCvect::getVec();
	typedef map<string,vector<vector<int> > >  ResMap;
	typedef ResMap::iterator ResIter;

	ResMap & ResDef=MyTop.getDef();
	vector<vector<int> > & Res=ResDef[y];
	MyRes=ResDef[y];
	Matrix co=Mt.getCO();
	Matrix oc=Mt.getOC();
	double v[DIM];
	v[XX]=co[XX][XX];v[YY]=co[YY][YY];v[ZZ]=co[ZZ][ZZ];
	double MinCO=*min_element(v,v+3)*0.5;  // Half the smallest axis is the cutoff distance for atom shifts
	Dvect Xref;

	// Apply boundary condition on every molecule as the default boundaries n .xtc files are atomic

	for(ResIter it=ResDef.begin(); it != ResDef.end();it++){
		vector<vector<int> > & Res0=it->second;
		for(unsigned int o=0;o<Res0.size();o++){
			Xref=xa[Res0[o][0]];
			for(unsigned int p=1;p<Res0[o].size();p++) {
				Dvect x1=Xref-xa[Res0[o][p]];
				Dvect xc1=co*x1;
				if(xc1.Norm() > MinCO) {
					Dvect tmp=*min_element(nxyz.begin(),nxyz.end(),CellComp(x1,co));
					for(int q=0; q < DIM;q++) xa[Res0[o][p]][q]=xa[Res0[o][p]][q]-tmp[q];
				}
			}
		}
	}
	// > Take the reference type molecules and eliminate boundary conditions --> begin

	// >>  Compute center of mass for each molecule


	vector<Dvect> xaa;
	for(unsigned int o=0; o< Res.size();o++){
		Dvect cm=0.0;
		for(unsigned int p=0; p< Res[o].size();p++){
			int pp=Res[o][p];
			cm+=xa[pp];
		}
		cm/=static_cast<double>(Res[o].size());
		xaa.push_back(cm);
	}


	// >> Obtain the center of mass for the cluster of the reference molecules

	Xref=xaa[0];
	for(unsigned n=1;n < xaa.size();n++){
		Dvect xb=xaa[n];
		Dvect tmp=*min_element(nxyz.begin(),nxyz.end(),CellComp(Xref,xb,co));
		xaa[n]=xaa[n]-tmp;
	}

	Dvect cma=0.0;
	for(unsigned n=0;n < xaa.size();n++){
		cma+=xaa[n];
	}
	cma/=static_cast<double> (xaa.size());
	// >>Shift all system by cma

	for(int n=0;n<nr;n++)
		for(int m=0;m<DIM;m++)
			xa[n][m]=xa[n][m]-cma[m];


	//> Apply molecular boundary conditions on the entire system
	for(ResIter it=ResDef.begin(); it != ResDef.end();it++){
		vector<vector<int> > & Res0=it->second;
		for(unsigned int o=0;o<Res0.size();o++){
			Dvect xcm=0.0;
			for(unsigned int p=0;p< Res0[o].size();p++) {
				xcm+=xa[Res0[o][p]];
			}

			xcm/=static_cast<double>(Res0[o].size());

			for(int m=0;m<DIM;m++) xcm[m]=rint(xcm[m]);
			for(unsigned int p=0;p<Res0[o].size();p++)
				for(int m=0;m<DIM;m++)
					xa[Res0[o][p]][m]=xa[Res0[o][p]][m]-xcm[m];

		}
	}

		//> Obtain new Cartesian coordinates from reduced xa's
	for(int i=0;i<nr;i++){
		for(int o=0;o<DIM;o++){
			xa[i][o]=xa[i][o]+0.5;
			x[i][o]=Mt.getCO()[o][XX]*xa[i][XX]+Mt.getCO()[o][YY]*xa[i][YY]+Mt.getCO()[o][ZZ]*xa[i][ZZ];
		}
	}
}

MAtoms::~MAtoms() {
	// TODO Auto-generated destructor stub
}

ostream & operator<<(ostream & fout, MAtoms & y){
	vector<Dvect> xc=vector<Dvect>(y.nr);
	if(1){
		for(int n=0;n<y.nr;n++){
			xc[n]=y.x[n];
		}
	} else{
		Matrix Rot0=y.InertiaRotCenter();
		Dvect cm=y.CMCenter();
		for(int n=0;n<y.nr;n++){
			Dvect xb=y.x[n];
			xb=xb-cm;
			xc[n]=Rot0*xb+cm;
		}
	}
	int ia=0;
	for(size_t i=0;i<y.CurrentPDB.size();i++){
		string Curr=y.CurrentPDB[i];
		if(Curr.find("ATOM") == 0 || Curr.find("HETATM") == 0 ) {
			string sub2=Curr.substr(16,4);
			sub2.erase(remove_if(sub2.begin(),sub2.end(),::isspace),sub2.end());
			if(!y.ResList0 || find(y.ResList0->begin(),y.ResList0->end(),sub2) != y.ResList0->end()) {
				fout << Curr.substr(0,29) ;
				fout << fixed << setw(8) << setprecision(3) << right << xc[ia][XX]/unit_nm;
				fout << fixed << setw(8) << setprecision(3) << right << xc[ia][YY]/unit_nm;
				fout << fixed << setw(8) << setprecision(3) << right << xc[ia][ZZ]/unit_nm;
				fout << Curr.substr(54,Curr.size()) << endl;
			}
			ia++;
		} else {
			fout << Curr <<endl;
			if(i==0){
				stringstream ss;
				ss<<y.time_c;
				fout << string("REMARK    SIMULATION TIME = "+ ss.str()) <<endl;
			}
		}
	}
	return fout;
}
Fstream & operator<<(Fstream & fout, MAtoms & y){
	FstreamC * foutC=dynamic_cast<FstreamC *> (&fout);
	t_fileio *xd = foutC->getfin();
	gmx_bool bOK=y.bOK_c;
	int step=y.step_c;
	real prec=y.prec_c;
	real time=y.time_c;
	matrix box;

	vector<Dvect> x0=vector<Dvect>(y.ResIndx0->size());
	for(size_t i=0;i < x0.size();i++){
		int ia=y.ResIndx0->at(i);
		x0.push_back(y.x[ia]);
	}
	int    natos=static_cast<int> (x0.size());
	rvec * x=new rvec [natos];
	for(unsigned int m=0;m<x0.size();m++)
		for(int o=0;o<DIM;o++) x[m][o]=x0[m][o];

	for(int i=0;i<DIM;i++)
			for(int j=0;j<DIM;j++)
				box[i][j]=y.Mt.getCO()[i][j];
	try{
		if(!write_xtc(xd,natos,step,time,box,x,prec)) throw string("Cannot write next frame. Abort");
	}catch(const string & s){
		cout << s << endl;
		exit(1);
	}
	return fout;
}



